import jieba
import pandas as pd
import matplotlib.pyplot as plt
import wordcloud
from imageio import imread
import numpy as np


raw = pd.read_table("金庸-射雕英雄传txt精校版.txt",names = ['txt'],  encoding ="GBK")
raw.head()

def h_head(tmpstr):
    return tmpstr[:1]

def h_mid(tmpstr):
    return tmpstr.find("回 ")
#取出进行判断的变量
raw['head'] = raw.txt.apply(h_head)
raw['mid'] = raw.txt.apply(h_mid)
raw['len'] = raw.txt.apply(len)#防止特殊情况，设置长度
raw.head()

chapnum = 0
for i in range(len(raw)):
    if raw['head'][i] == "第" and raw['mid'][i] > 0 and raw['len'][i] < 30:
        chapnum += 1
    if chapnum >= 40 and raw['txt'][i] == "附录一：成吉思汗家族":
        chapnum = 0
    raw.loc[i, 'chap'] = chapnum

del raw['head']
del raw['mid']
del raw['len']

tmpchap = raw[raw['chap'] == 1].copy()
tmpchap.reset_index(drop=True, inplace=True)
tmpchap['paraidx'] = tmpchap.index
raw1=tmpchap.txt[1]

dict = '金庸小说词库.txt'
jieba.load_userdict(dict)
res1 = jieba.lcut(raw1)

myfont = r'C:\Windows\Fonts\simkai.ttf'

def m_cut(intxt):
    return [ w for w in jieba.cut(intxt) if w not in stoplist and len(w) > 1]
stoplist = list(pd.read_table('停用词.txt', names = ['w'],
                            encoding = 'utf-8', engine='python').w)

rawchap = [ " ".join(m_cut(w)) for w in tmpchap.txt.iloc]


udobj = wordcloud.WordCloud(font_path = myfont,mode = "RGBA", background_color = None).\
    generate(' '.join(rawchap))
plt.imshow(udobj)
plt.axis("off")
plt.show()

tmpdf_hum = pd.read_table('人名.txt',names = ['w'], encoding='GBK')
human_names = [ w for w in res1 if w in list(tmpdf_hum.w) ]
tmpdf_pla = pd.read_table('金庸地名.txt',names = ['w'], encoding='GBK')
place_names = [ w for w in res1 if w in list(tmpdf_pla.w) ]

import matplotlib.pyplot as plt
cloudhum = wordcloud.WordCloud(font_path = myfont,mode = "RGBA", background_color = None).\
    generate(' '.join(human_names) )
plt.imshow(cloudhum)
plt.axis("off")
plt.show()

cloudpla = wordcloud.WordCloud(font_path = myfont,mode = "RGBA", background_color = None).\
    generate(' '.join(place_names) )
plt.imshow(cloudpla)
plt.axis("off")
plt.show()

nam=str(human_names).join(str(place_names))

cloudobj = wordcloud.WordCloud(font_path = myfont,mode = "RGBA", background_color = None)\
    .generate(' '.join(nam))

plt.imshow(cloudobj)
plt.axis("off")
plt.show()

from wordcloud import get_single_color_func

class GroupedColorFunc(object):

    def __init__(self, color_to_words, default_color):
        self.color_func_to_words = [
            (get_single_color_func(color), set(words))
            for (color, words) in color_to_words.items()]

        self.default_color_func = get_single_color_func(default_color)

    def get_color_func(self, word):
        try:
            color_func = next(
                color_func for (color_func, words) in self.color_func_to_words
                if word in words)
        except StopIteration:
            color_func = self.default_color_func

        return color_func

    def __call__(self, word, **kwargs):
        return self.get_color_func(word)(word, **kwargs)

color_to_words = {'blue': human_names,'red': place_names}
default_color = 'black'
grouped_color_func = GroupedColorFunc(color_to_words, default_color)
udobj.recolor(color_func=grouped_color_func)
plt.imshow(udobj)
plt.axis("off")
plt.show()

imgobj = imread("射雕英雄传.jpg")
image_colors = wordcloud.ImageColorGenerator(np.array(imgobj))
udobj.recolor(color_func=image_colors)

plt.imshow(udobj)
plt.axis("off")
plt.show()

def h_cut(intxt):
    return [ w for w in jieba.cut(intxt) if w not in stoplist and len(w) > 1 ]
udobj = wordcloud.WordCloud(font_path = myfont,mask = imread("绿色图片.jpg"),mode = "RGBA",
                    background_color = None).generate(' '.join(res1))

plt.imshow(udobj)
plt.axis("off")
plt.show()